import numpy as np
import matplotlib.pyplot as plt

def read_loss_data(filename):
    """Read loss data from a file with 'iteration loss' format."""
    iterations = []
    losses = []
    with open(filename, 'r') as file:
        for line in file:
            try:
                # Parse line as "iteration loss" format
                parts = line.strip().split()
                if len(parts) == 2:
                    iteration = float(parts[0])
                    loss = float(parts[1])
                    iterations.append(iteration)
                    losses.append(loss)
            except ValueError:
                # Skip lines that can't be parsed
                continue
    return iterations, losses

# File paths
file_128 = 'loss_prims_128.txt'
file_4096 = 'loss_prims_4096.txt'

# Read data
iterations_128, losses_128 = read_loss_data(file_128)
iterations_4096, losses_4096 = read_loss_data(file_4096)

# Create the plot
plt.figure(figsize=(10, 6))
plt.plot(iterations_128, losses_128, 'b-', linewidth=2, label='Primitives: 128')
plt.plot(iterations_4096, losses_4096, 'r-', linewidth=2, label='Primitives: 4096')

# Set axis labels and title
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.title('Loss Comparison: 128 vs 4096 Primitives')

# Add grid and legend
plt.grid(True, linestyle='--', alpha=0.7)
plt.legend()

# Improve visual appearance
plt.tight_layout()

# Optional: Set logarithmic scale if loss values span multiple orders of magnitude
# plt.yscale('log')

# Save the plot
plt.savefig('loss_comparison.png', dpi=300)

# Display the plot
plt.show()

print("Plot saved as 'loss_comparison.png'")
